import os
import cv2
 
import numpy as np
from skimage.feature import hog
from PIL import ImageFont, ImageDraw, Image
class ImageClassifier:
    def __init__(self):
        pass
        

    def load_images_from_folder(self,folder):
        self.folder = folder
        images = []
        labels = []
        for label in os.listdir(self.folder):
            label_folder = os.path.join(self.folder, label)
            if not os.path.isdir(label_folder):
                continue
            for filename in os.listdir(label_folder):
                img_path = os.path.join(label_folder, filename)
                img = cv2.imread(img_path)
                if img is not None:
                    img = cv2.resize(img, (64, 64))  # 调整图像大小
                    images.append(img)
                    labels.append(label)
        return images, labels

    def extract_hog_features(self, input):
        def process_image(img):
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图像
            fd, _ = hog(gray, pixels_per_cell=(8, 8), cells_per_block=(2, 2), visualize=True, channel_axis=None)
            return fd

        if isinstance(input, list):  # 如果输入是一个列表
            hog_features = []
            for img in input:
                fd = process_image(img)
                hog_features.append(fd)
            return np.array(hog_features)
        elif isinstance(input, np.ndarray):
                return process_image(input)

    def extract_color_features(self, images):
        def process_image(img):
            hist_r = cv2.calcHist([img], [0], None, [256], [0, 256])
            hist_g = cv2.calcHist([img], [1], None, [256], [0, 256])
            hist_b = cv2.calcHist([img], [2], None, [256], [0, 256])
            hist = np.concatenate((hist_r, hist_g, hist_b)).flatten()
            hist = hist / np.sum(hist)
            return hist
        if isinstance(images, list):  # 如果输入是一个列表
            color_features = []
            for img in images:
                hist = process_image(img)
                color_features.append(hist)
            return np.array(color_features)
        elif isinstance(images, np.ndarray):
            return process_image(images)
    
    def drawChinese(self,text,x,y,size,r, g, b, a,img,font):
        if font == None:
            frame = cv2.putText(img, text, (x, y),cv2.FONT_HERSHEY_SIMPLEX, size, (r, g, b), 2)
        else:
            font = ImageFont.truetype(font, size)
            img_pil = Image.fromarray(img)
            draw = ImageDraw.Draw(img_pil)
            draw.text((x,y), text, font=font, fill=(b, g, r, a))
            frame = np.array(img_pil)
        return frame